import tensorflow as tf
from tensorflow.keras.datasets import mnist # type: ignore
from tensorflow.keras.models import Model
import numpy as np
import matplotlib.pyplot as plt
from LeNet5 import LeNet5

# 加载MNIST数据集
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# 调整数据形状以适应LeNet5的输入(28*28 -> 32*32)
train_images = np.pad(train_images, ((0, 0), (2, 2), (2, 2)), 'constant', constant_values=0)
test_images = np.pad(test_images, ((0, 0), (2, 2), (2, 2)), 'constant', constant_values=0)
train_images = train_images.reshape((60000, 32, 32, 1))
test_images = test_images.reshape((10000, 32, 32, 1))

# 归一化处理
train_images = train_images.astype('float32') / 255
test_images = test_images.astype('float32') / 255

# 创建LeNet - 5模型实例
model = LeNet5()

# 编译模型
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 训练模型
history = model.fit(train_images, train_labels, epochs=10, batch_size=128, validation_data=(test_images, test_labels))

# 可视化训练过程
# 绘制训练和验证损失曲线
plt.plot(history.history['loss'], label='Training Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

# 绘制训练和验证准确率曲线
plt.plot(history.history['accuracy'], label='Training Accuracy')
plt.plot(history.history['val_accuracy'], label='Validation Accuracy')
plt.title('Training and Validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

# 评估模型
test_loss, test_acc = model.evaluate(test_images, test_labels)
print(f"Test accuracy: {test_acc}")

# 保存模型
model.save('./LeNet5/lenet5_model.h5')
print("Model saved as lenet5_model.h5")
